#include "octree.h"
#include <boost/algorithm/string/predicate.hpp>
#include <fstream>
#include <iostream>

namespace {
double DisPtToLine(const Eigen::Vector3d& pt, const Eigen::Vector3d& ptStart,
                   const Eigen::Vector3d& ptEnd) {
  double dLength = (ptStart - ptEnd).norm();
  if (dLength < 1e-4) {
    return (ptStart - pt).norm();
  }

  Eigen::Vector3d e1 = ptEnd - ptStart;
  Eigen::Vector3d e2 = pt - ptStart;
  Eigen::Vector3d n = e1.cross(e2);
  return n.norm() / dLength;
};
}

void LeafNode::AddRecurPoint(const las_Point& las_pt, double resolution) {
  points_.emplace_back(las_pt);
}

bool LeafNode::IsLeaf() const { return true; }

void LeafNode::Archive(std::ostream& ofs) {
  ofs.write(reinterpret_cast<const char*>(&center_.x()), sizeof(float));
  ofs.write(reinterpret_cast<const char*>(&center_.y()), sizeof(float));
  ofs.write(reinterpret_cast<const char*>(&center_.z()), sizeof(float));
  ofs.write(reinterpret_cast<const char*>(&radius_), sizeof(float));

  int num = points_.size();
  ofs.write(reinterpret_cast<const char*>(&num), sizeof(int));
  for (int i = 0; i < num; i++) {
    las_Point pt = points_[i];
    ofs.write(reinterpret_cast<const char*>(&pt.point_x), sizeof(float));
    ofs.write(reinterpret_cast<const char*>(&pt.point_y), sizeof(float));
    ofs.write(reinterpret_cast<const char*>(&pt.point_z), sizeof(float));
    ofs.write(reinterpret_cast<const char*>(&pt.color_r), sizeof(uint8_t));
    ofs.write(reinterpret_cast<const char*>(&pt.color_g), sizeof(uint8_t));
    ofs.write(reinterpret_cast<const char*>(&pt.color_b), sizeof(uint8_t));
  }
}

void LeafNode::Archive(std::istream& ifs) {
  ifs.read(reinterpret_cast<char*>(&center_.x()), sizeof(float));
  ifs.read(reinterpret_cast<char*>(&center_.y()), sizeof(float));
  ifs.read(reinterpret_cast<char*>(&center_.z()), sizeof(float));
  ifs.read(reinterpret_cast<char*>(&radius_), sizeof(float));

  int num = 0;
  ifs.read(reinterpret_cast<char*>(&num), sizeof(int));
  points_.resize(num);

  for (int i = 0; i < num; i++) {
    las_Point pt;
    ifs.read(reinterpret_cast<char*>(&pt.point_x), sizeof(float));
    ifs.read(reinterpret_cast<char*>(&pt.point_y), sizeof(float));
    ifs.read(reinterpret_cast<char*>(&pt.point_z), sizeof(float));
    ifs.read(reinterpret_cast<char*>(&pt.color_r), sizeof(uint8_t));
    ifs.read(reinterpret_cast<char*>(&pt.color_g), sizeof(uint8_t));
    ifs.read(reinterpret_cast<char*>(&pt.color_b), sizeof(uint8_t));
    points_[i] = pt;
  }

}

const std::vector<las_Point>& LeafNode::GetPoints() const { return points_; }

void LeafNode::CyberSimplify(int cyber_num) {
  float cyber_length = radius_ / cyber_num * 2;
  std::map<
      std::tuple<int, int, int>,
      std::map<std::tuple<uint8_t, uint8_t, uint8_t>, std::vector<las_Point>>>
      map_cyber_points;
  for (las_Point pt : points_) {
    Eigen::Vector3f local(pt.point_x - center_.x(), pt.point_y - center_.y(),
                          pt.point_z - center_.z());
    int row = local.x() / cyber_length;
    int col = local.y() / cyber_length;
    int hgt = local.z() / cyber_length;
    std::tuple<int, int, int> idx = std::make_tuple(row, col, hgt);
    std::tuple<uint8_t, uint8_t, uint8_t> color =
        std::make_tuple(pt.color_r, pt.color_g, pt.color_b);
    map_cyber_points[idx][color].push_back(pt);
  }

  std::vector<las_Point> simple_pts;
  for (auto cyber_mit = map_cyber_points.begin();
       cyber_mit != map_cyber_points.end(); cyber_mit++) {
    std::map<std::tuple<uint8_t, uint8_t, uint8_t>, std::vector<las_Point>>
        map_label_points = cyber_mit->second;
    for (auto mit = map_label_points.begin(); mit != map_label_points.end();
         mit++) {
      std::vector<las_Point> las_points = mit->second;
      if (!las_points.empty()) {
        simple_pts.emplace_back(las_points[0]);
      }
    }
  }
  points_.swap(simple_pts);
}

BranchNode::BranchNode() {
  for (int i = 0; i < 8; i++) {
    children_[i] = nullptr;
  }
}

BranchNode::~BranchNode() {
  for (int i = 0; i < 8; i++) {
    if (children_[i]) {
      delete children_[i];
      children_[i] = nullptr;
    }
  }
}

void BranchNode::AddRecurPoint(const las_Point& las_pt, double resolution) {
  int child_idx = GetChildIdx(las_pt.point_x, las_pt.point_y, las_pt.point_z);
  octreeNode* child = CreateOrGetChild(child_idx, resolution);
  if (child) {
    child->AddRecurPoint(las_pt, resolution);
  }
}

octreeNode* BranchNode::CreateOrGetChild(int idx, double resolution) {
  if (idx == -1) {
    return nullptr;
  }

  if (children_[idx] != nullptr) {
    return children_[idx];
  }

  Eigen::Vector3f center;
  double radius;
  GetChildExtent(idx, center, radius);

  if (radius <= resolution / 2.0) {
    LeafNode* child = new LeafNode;
    child->SetExtent(center, radius);
    children_[idx] = child;
    return children_[idx];
  }

  BranchNode* child = new BranchNode;
  child->SetExtent(center, radius);
  children_[idx] = child;
  return children_[idx];
}

bool BranchNode::IsLeaf() const { return false; }

void BranchNode::GetChildExtent(int idx, Eigen::Vector3f& center,
                                double& radius) {
  int tx = (idx & 0x01) ? 1 : 0;
  int ty = (idx & 0x02) ? 1 : 0;
  int tz = (idx & 0x04) ? 1 : 0;

  center = center_;
  radius = radius_ / 2.0;

  if (tx == 1) {
    center.x() += radius;
  } else {
    center.x() -= radius;
  }

  if (ty == 1) {
    center.y() += radius;
  } else {
    center.y() -= radius;
  }

  if (tz == 1) {
    center.z() += radius;
  } else {
    center.z() -= radius;
  }
}

int BranchNode::GetChildIdx(float x, float y, float z) {
  if (x < center_.x() - radius_ || y < center_.y() - radius_ ||
      z < center_.z() - radius_ || x > center_.x() + radius_ ||
      y > center_.y() + radius_ || z > center_.z() + radius_) {
    return -1;
  }

  int tx = (x <= center_.x()) ? 0 : 1;
  int ty = (y <= center_.y()) ? 0 : 1;
  int tz = (z <= center_.z()) ? 0 : 1;

  int idx = tx + (ty << 1) + (tz << 2);
  return idx;
}

void BranchNode::Archive(std::ostream& ofs) {
  ofs.write(reinterpret_cast<const char*>(&center_.x()), sizeof(float));
  ofs.write(reinterpret_cast<const char*>(&center_.y()), sizeof(float));
  ofs.write(reinterpret_cast<const char*>(&center_.z()), sizeof(float));
  ofs.write(reinterpret_cast<const char*>(&radius_), sizeof(float));

  int child_num = 0;
  for (int i = 0; i < 8; i++) {
    octreeNode* child = children_[i];
    if (child) {
      child_num++;
    }
  }
  ofs.write(reinterpret_cast<const char*>(&child_num), sizeof(int));
  for (int i = 0; i < 8; i++) {
    octreeNode* child = children_[i];
    if (child) {
      bool is_child_leaf = child->IsLeaf();
      ofs.write(reinterpret_cast<const char*>(&i), sizeof(int));
      ofs.write(reinterpret_cast<const char*>(&is_child_leaf), sizeof(bool));
      child->Archive(ofs);
    }
  }
}

void BranchNode::Archive(std::istream& ifs) {
  ifs.read(reinterpret_cast<char*>(&center_.x()), sizeof(float));
  ifs.read(reinterpret_cast<char*>(&center_.y()), sizeof(float));
  ifs.read(reinterpret_cast<char*>(&center_.z()), sizeof(float));
  ifs.read(reinterpret_cast<char*>(&radius_), sizeof(float));

  int child_num = 0;
  ifs.read(reinterpret_cast<char*>(&child_num), sizeof(int));
  for (int i = 0; i < child_num; i++) {
    int child_idx = 0;
    bool is_child_leaf = false;
    ifs.read(reinterpret_cast<char*>(&child_idx), sizeof(int));
    ifs.read(reinterpret_cast<char*>(&is_child_leaf), sizeof(bool));

    if (is_child_leaf) {
      LeafNode* node = new LeafNode;
      node->Archive(ifs);
      children_[child_idx] = node;
    }
    else {
      BranchNode* node = new BranchNode;
      node->Archive(ifs);
      children_[child_idx] = node;
    }
  }
}

octreeNode::octreeNode() : radius_(0), center_(0, 0, 0) {}

octreeNode::~octreeNode() {}

const Eigen::Vector3f& octreeNode::GetCenter() { return center_; }

float octreeNode::GetRadius() { return radius_; }

void octreeNode::SetExtent(const Eigen::Vector3f& center, double radius) {
  center_ = center;
  radius_ = radius;
}

void octreeNode::GetExtent(Eigen::Vector3f& min, Eigen::Vector3f& max) {
  min = center_ - Eigen::Vector3f(radius_, radius_, radius_);
  max = center_ + Eigen::Vector3f(radius_, radius_, radius_);
}

octree::octree()
    : resolution_(0.2),
      block_size_(50.0),
      cyber_num_(4),
      offset_(0.0, 0.0, 0.0) {
  min_.x() = std::numeric_limits<double>::max();
  min_.y() = std::numeric_limits<double>::max();
  min_.z() = std::numeric_limits<double>::max();
  max_.x() = std::numeric_limits<double>::min();
  max_.y() = std::numeric_limits<double>::min();
  max_.z() = std::numeric_limits<double>::min();
}

octree::~octree() {
 
}

void octree::SetResolution(const double& resolution) {
  resolution_ = resolution;
}

double octree::GetResolution() const { return resolution_; }

void octree::SetCyberNum(int cyber_num) { cyber_num_ = cyber_num; }

int octree::GetCyberNum() const { return cyber_num_; }

void octree::AddPoint(const las_Point& las_pt) {
  if (las_pt.point_x > max_.x()) max_.x() = las_pt.point_x;
  if (las_pt.point_x < min_.x()) min_.x() = las_pt.point_x;
  if (las_pt.point_y > max_.y()) max_.y() = las_pt.point_y;
  if (las_pt.point_y < min_.y()) min_.y() = las_pt.point_y;
  if (las_pt.point_z > max_.z()) max_.z() = las_pt.point_z;
  if (las_pt.point_z < min_.z()) min_.z() = las_pt.point_z;

  int idx = las_pt.point_x / block_size_;
  int idy = las_pt.point_y / block_size_;
  int idz = las_pt.point_z / block_size_;
  if (idx <= 0 && las_pt.point_x < 0) idx--;
  if (idy <= 0 && las_pt.point_y < 0) idy--;
  if (idz <= 0 && las_pt.point_z < 0) idz--;

  std::tuple<int, int, int> id = std::make_tuple(idx, idy, idz);
  Eigen::Vector3f center = Eigen::Vector3f::Zero();
  center.x() = (idx + 0.5) * block_size_;
  center.y() = (idy + 0.5) * block_size_;
  center.z() = (idz + 0.5) * block_size_;
  double radius = block_size_ / 2.0;

  if (roots_.find(id) == roots_.end()) {
    roots_[id] = std::make_shared<BranchNode>();
    roots_[id]->SetExtent(center, radius);
  }
  roots_[id]->AddRecurPoint(las_pt, resolution_);
}

std::vector<octreeNode*> octree::GetLeaves() const {
  std::vector<octreeNode*> leaves;
  std::function<void(octreeNode*)> getchild = [&](octreeNode* node) -> void {
    if (node) {
      if (node->IsLeaf()) {
        leaves.push_back(node);
      } else {
        BranchNode* branch = static_cast<BranchNode*>(node);
        octreeNode** children = branch->children_;
        for (int i = 0; i < 8; i++) {
          getchild(children[i]);
        }
      }
    }
  };
  for (auto mit = roots_.begin(); mit != roots_.end(); mit++) {
    getchild(mit->second.get());
  }
  return leaves;
}

Eigen::Vector3d octree::GetOffset() const { return offset_; }

void octree::Simplify(int min_points_leave) {
  std::function<void(octreeNode*)> simply = [&](octreeNode* node) -> void {
    if (node->IsLeaf()) {
      LeafNode* leaf_node = static_cast<LeafNode*>(node);
      if (leaf_node->points_.size() > min_points_leave) {
        leaf_node->CyberSimplify(cyber_num_);
      } else {
        std::vector<las_Point>().swap(leaf_node->points_);
      }
    } else {
      BranchNode* branch = static_cast<BranchNode*>(node);
      octreeNode** children = branch->children_;
      for (int i = 0; i < 8; i++) {
        if (children[i]) {
          simply(children[i]);
        }
      }
    }
  };
  for (auto mit = roots_.begin(); mit != roots_.end(); mit++) {
    simply(mit->second.get());
  }
}

std::vector<las_Point> octree::GetPoints() const {
  std::vector<octreeNode*> leaves = GetLeaves();

  std::vector<las_Point> total_points;
  for (octreeNode* leaf : leaves) {
    LeafNode* leaf_node = static_cast<LeafNode*>(leaf);
    std::vector<las_Point> points = leaf_node->GetPoints();
    for (const las_Point& center : points) {
      las_Point simple_pt;
      simple_pt.point_x = center.point_x;
      simple_pt.point_y = center.point_y;
      simple_pt.point_z = center.point_z;
      simple_pt.color_r = center.color_r;
      simple_pt.color_g = center.color_g;
      simple_pt.color_b = center.color_b;
      total_points.emplace_back(simple_pt);
    }
  }
  return total_points;
}

bool octree::IntersectWithLine(const Eigen::Vector3d& start_pt,
                               const Eigen::Vector3d& end_pt,
                               Eigen::Vector3d& ptHit) const {
  Eigen::Vector3d s2e = end_pt - start_pt;
  s2e.normalize();

  std::vector<std::pair<double, Eigen::Vector3d>> distances;
  std::function<void(octreeNode*)> intersect = [&](octreeNode* node) -> bool {
    Eigen::Vector3d center = node->GetCenter().cast<double>() + offset_;
    double dis = DisPtToLine(center, start_pt, end_pt);
    if (dis > 1.732 * node->GetRadius()) {
      return false;
    }

    if (node->IsLeaf()) {
      LeafNode* leaveNode = static_cast<LeafNode*>(node);
      std::vector<las_Point> points = leaveNode->GetPoints();
      for (las_Point pt : points) {
        Eigen::Vector3d point(pt.point_x + offset_.x(),
                              pt.point_y + offset_.y(),
                              pt.point_z + offset_.z());
        double dis = DisPtToLine(point, start_pt, end_pt);
        distances.emplace_back(dis, point);
      }
    } else {
      BranchNode* branch = static_cast<BranchNode*>(node);
      octreeNode** children = branch->children_;
      for (int i = 0; i < 8; i++) {
        if (children[i]) {
          intersect(children[i]);
        }
      }
    }
    return true;
  };

  for (auto mit = roots_.begin(); mit != roots_.end(); mit++) {
    std::shared_ptr<BranchNode> root = mit->second;
    Eigen::Vector3d center = root->GetCenter().cast<double>() + offset_;
    double dis = DisPtToLine(center, start_pt, end_pt);
    if (dis > 1.732 * root->GetRadius()) {
      continue;
    }

    intersect(root.get());
  }

  if (distances.empty()) {
    return false;
  }
  std::sort(distances.begin(), distances.end(),
            [=](const std::pair<double, Eigen::Vector3d>& lhs,
                const std::pair<double, Eigen::Vector3d>& rhs) -> bool {
              return lhs.first < rhs.first;
            });
  ptHit = distances[0].second;
  return true;
}

bool octree::ReadPlyFile(const std::string& filename) {
  std::ifstream ifs(filename);
  if (!ifs) {
    return false;
  }
  int points_num = 0;
  Eigen::Vector3d offset = Eigen::Vector3d::Zero();

  std::string buffer;
  while (getline(ifs, buffer)) {
    if (boost::starts_with(buffer, "element vertex")) {
      buffer = buffer.substr(strlen("element vertex"));
      points_num = std::stoi(buffer);
      continue;
    } 
    
    std::size_t found = buffer.find("offset_x");
    if (found != std::string::npos) {
      buffer = buffer.substr(found + strlen("offset_x"));
      offset.x() = std::stod(buffer);
      continue;
    } 

    found = buffer.find("offset_y");
    if (found != std::string::npos) {
      buffer = buffer.substr(found + strlen("offset_y"));
      offset.y() = std::stod(buffer);
      continue;
    }

    found = buffer.find("offset_z");
    if (found != std::string::npos) {
      buffer = buffer.substr(found + strlen("offset_z"));
      offset.z() = std::stod(buffer);
      continue;
    }

    if (boost::starts_with(buffer, "end_header")) {
      break;
    }
  }

  //----------------------
  // read vertices
  for (int i = 0; i < points_num; i++) {
    getline(ifs, buffer);
    std::stringstream sstream(buffer);

    float x, y, z;
    int r, g, b;
    sstream >> x >> y >> z >> r >> g >> b;

    las_Point point;
    point.point_x = x;
    point.point_y = y;
    point.point_z = z;
    point.color_r = r;
    point.color_g = g;
    point.color_b = b;
    AddPoint(point);
  }
  offset_ = offset;

  ifs.close();
  return true;
}

bool octree::WriteOctreeFile(const std::string& filename) {
  std::ofstream ofs(filename, std::ios::binary);
  ofs.write(reinterpret_cast<const char*>(&resolution_), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&cyber_num_), sizeof(int));
  ofs.write(reinterpret_cast<const char*>(&offset_.x()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&offset_.y()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&offset_.z()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&min_.x()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&min_.y()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&min_.z()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&max_.x()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&max_.y()), sizeof(double));
  ofs.write(reinterpret_cast<const char*>(&max_.z()), sizeof(double));

  int size = roots_.size();
  ofs.write(reinterpret_cast<const char*>(&size), sizeof(int));
  for (auto mit = roots_.begin(); mit != roots_.end(); mit++) {
    int idx, idy, idz;
    std::tie(idx, idy, idz) = mit->first;
    ofs.write(reinterpret_cast<const char*>(&idx), sizeof(int));
    ofs.write(reinterpret_cast<const char*>(&idy), sizeof(int));
    ofs.write(reinterpret_cast<const char*>(&idz), sizeof(int));
    mit->second->Archive(ofs);
  }
  return true;
}

bool octree::ReadOctreeFile(const std::string& filename) {
  std::ifstream ifs(filename, std::ios::binary);
  ifs.read(reinterpret_cast<char*>(&resolution_), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&cyber_num_), sizeof(int));
  ifs.read(reinterpret_cast<char*>(&offset_.x()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&offset_.y()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&offset_.z()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&min_.x()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&min_.y()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&min_.z()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&max_.x()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&max_.y()), sizeof(double));
  ifs.read(reinterpret_cast<char*>(&max_.z()), sizeof(double));
  
  int size = 0;
  ifs.read(reinterpret_cast<char*>(&size), sizeof(int));
  for (int i = 0; i < size; i++) {
    int idx, idy, idz;
    ifs.read(reinterpret_cast<char*>(&idx), sizeof(int));
    ifs.read(reinterpret_cast<char*>(&idy), sizeof(int));
    ifs.read(reinterpret_cast<char*>(&idz), sizeof(int));

    std::tuple<int, int, int> id = std::make_tuple(idx, idy, idz);
    roots_[id] = std::make_shared<BranchNode>();
    roots_[id]->Archive(ifs);
  }
  return true;
}
